#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <omp.h>

//Doing just keccak[1600] for now...

#define STATE_ROWS 5

/*const int rot_const[5][5] = {
	{0, 36, 3, 41, 18},
	{1, 44, 10, 45, 2},
	{62, 6, 43, 15, 61},
	{28, 55, 25, 21, 56},
	{27, 20, 39, 8, 14}
};*/

/*const int rot_const[25] = {
	0, 36, 3, 41, 18,
	1, 44, 10, 45, 2,
	62, 6, 43, 15, 61,
	28, 55, 25, 21, 56,
	27, 20, 39, 8, 14
};*/

const int rot_const[5][5] = {
	{0, 36, 3, 41, 18},
	{1, 44, 10, 45, 2},
	{62, 6, 43, 15, 61},
	{28, 55, 25, 21, 56},
	{27, 20, 39, 8, 14}
};

const int B_indices[25] = {
		0,8,11,19,22,
		2,5,13,16,24,
		4,7,10,18,21,
		1,9,12,15,23,
		3,6,14,17,20
	};
	
const uint64_t RC[24] = { 0x0000000000000001,0x0000000000008082,0x800000000000808A,0x8000000080008000,0x000000000000808B,0x0000000080000001,
					0x8000000080008081,0x8000000000008009,0x000000000000008A,0x0000000000000088,0x0000000080008009,0x000000008000000A,
					0x000000008000808B,0x800000000000008B,0x8000000000008089,0x8000000000008003,0x8000000000008002,0x8000000000000080,
					0x000000000000800A,0x800000008000000A,0x8000000080008081,0x8000000000008080,0x0000000080000001,0x8000000080008008
				};

inline uint64_t rot(uint64_t in, int n)
{
	return (in << n) | (in >> (64 - n));
}

inline uint64_t reverse(uint64_t in)
{
	return (in << 56) | ((in & 0xFF00) << 40) | ((in & 0xFF0000) << 24) | ((in & 0xFF000000) << 8) | ((in & 0xFF00000000) >> 8) | ((in & 0xFF0000000000) >> 24) | ((in & 0xFF000000000000) >> 40) | ((in & 0xFF00000000000000) >> 56);
}

int main_orig(int argc, char* argv[])
{
	uint8_t *init_message;
	uint8_t *padded_message;
	uint64_t* blocks;
	uint64_t message_length, padded_length;
	FILE *infile;
	struct stat buf;
	int i, bytes_to_pad, x, y, num_threads, j;
	int count = 0;
	char output[25*16+1] = "";
	char temp[17];
	uint64_t S[STATE_ROWS*STATE_ROWS] = {0};
	uint64_t B[STATE_ROWS*STATE_ROWS],C[STATE_ROWS],D[STATE_ROWS];
	double temp_time, profile_time = 0;
	double start;
	if(argc != 3)
	{
		printf("Usage: ./keccak <filename> <num_threads>\n");
		exit(1);
	}
	
	infile = fopen(argv[1], "rb");
	num_threads = atoi(argv[2]);
	omp_set_num_threads(num_threads);
	
	if(infile == NULL)
	{
		printf("Could not open input file %s.\n",argv[1]);
		exit(1);
	}
			
	fstat(fileno(infile), &buf);
	
	message_length = buf.st_size;
	//message_length = 1;

	//init_message = (uint8_t*)malloc(message_length*sizeof(uint8_t));
	bytes_to_pad = 72 - (message_length % 72);
	padded_length = message_length + bytes_to_pad;
	padded_message = (uint8_t*)calloc(padded_length,sizeof(uint8_t));
	//temp_time = omp_get_wtime();
	fread((void*)(padded_message),1,message_length,infile);
	//profile_time = omp_get_wtime() - temp_time;
	
	//printf("Message length = %d\nBytes to add = %d\nPadded length = %d\n",message_length,bytes_to_pad,padded_length);
	
	blocks = (uint64_t*)malloc(padded_length);
	
	//memcpy(padded_message,init_message,message_length);			
	padded_message[message_length] = 1;
	padded_message[padded_length-1] |= 0x80;
	memcpy(blocks,padded_message,padded_length);
	start = omp_get_wtime();
	#pragma omp parallel private(i,j)
	{
		for(i = 0; i < padded_length/8; i+=9)
		{
			#pragma omp for
			for(x = 0; x < 9; x++)
			{
				S[(5*x)%24] ^= blocks[x+i];
			}
			
			for(j = 0; j < 24; j++)
			{
				#pragma omp for
				for(x = 0; x < STATE_ROWS; x++)
				{
					C[x] = S[x*STATE_ROWS] ^ S[x*STATE_ROWS+1] ^ S[x*STATE_ROWS+2] ^ S[x*STATE_ROWS+3] ^ S[x*STATE_ROWS+4];
				}
				
				#pragma omp for
				for(x = 0; x < STATE_ROWS; x++)
				{
					D[x] = C[(x+4)%STATE_ROWS] ^ rot(C[(x+1)%STATE_ROWS],1);
				}
				
				//ρ and π steps
				#pragma omp for
				for(x = 0; x < STATE_ROWS*STATE_ROWS; x++)
				{
					S[x] = S[x] ^ D[x/5];
					B[B_indices[x]] = rot(S[x], rot_const[x]);
				}
				
				//χ step
				
				/*#pragma omp master
				{
				temp_time = omp_get_wtime();
				}*/
				#pragma omp for
				/*for(x = 0; x < 25; x++)
				{
					S[x] = B[x] ^ ((~B[(x+5)%25]) & B[(x+10)%25]);
				}*/
				//#pragma omp parallel for
				for(x = 0; x < 5; x++)
				{
					for(y = 0; y < 5; y++)
					{
						S[x*STATE_ROWS+y] = B[x*STATE_ROWS+y] ^ ((~B[((x+1)%STATE_ROWS)*STATE_ROWS+y]) & B[((x+2)%STATE_ROWS)*STATE_ROWS+y]);
					}
				}
				
				//ι step
				#pragma omp master
				{
					//profile_time += omp_get_wtime() - temp_time;
					S[0] = S[0] ^ RC[j];
				}
			}
		}
	
		//#pragma omp barrier
		
		#pragma omp for private(temp)
		for(x = 0; x < 8; x++)
		{
			sprintf(temp,"%.16lX",reverse(S[(5*x)%24]));
			memcpy(output+x*16,temp,16);
		}
	}
	start = omp_get_wtime()-start;
	//keccak_func(S);
	
	printf("%s %s\n",output,argv[1]);
	
	printf("Time: %lf\nProfiled part: %lf\n",start,profile_time);
	return 0;
}

int main(int argc, char* argv[])
{
	uint8_t *init_message;
	uint8_t *padded_message;
	uint64_t* blocks;
	uint64_t message_length, padded_length;
	FILE *infile;
	struct stat buf;
	int i, bytes_to_pad, x, y, j;
	int count = 0;
	int blocks_read = 0,my_blocks_read,num_threads;
	char output[25*16+1] = "";
	char temp[17];
	uint8_t temp_read[72] = {0};
	uint64_t B[STATE_ROWS*STATE_ROWS],C[STATE_ROWS],D[STATE_ROWS];
	uint64_t S[STATE_ROWS*STATE_ROWS] = {0};
	uint64_t RC[24] = { 0x0000000000000001,0x0000000000008082,0x800000000000808A,0x8000000080008000,0x000000000000808B,0x0000000080000001,
						0x8000000080008081,0x8000000000008009,0x000000000000008A,0x0000000000000088,0x0000000080008009,0x000000008000000A,
						0x000000008000808B,0x800000000000008B,0x8000000000008089,0x8000000000008003,0x8000000000008002,0x8000000000000080,
						0x000000000000800A,0x800000008000000A,0x8000000080008081,0x8000000000008080,0x0000000080000001,0x8000000080008008
					};
	double temp_time, profile_time = 0;
	double start;
	if(argc != 3)
	{
		printf("Usage: ./keccak <filename> <num_threads>\n");
		exit(1);
	}
	
	infile = fopen(argv[1], "rb");
	num_threads = atoi(argv[2]);
	omp_set_num_threads(num_threads);
	
	if(infile == NULL)
	{
		printf("Could not open input file %s.\n",argv[1]);
		exit(1);
	}
	
	fstat(fileno(infile), &buf);
	
	message_length = buf.st_size;
	//message_length = 1;

	//init_message = (uint8_t*)malloc(message_length*sizeof(uint8_t));
	bytes_to_pad = 72 - (message_length % 72);
	padded_length = message_length + bytes_to_pad;
	padded_message = (uint8_t*)calloc(padded_length,sizeof(uint8_t));
	padded_message[message_length] = 1;
	padded_message[padded_length-1] |= 0x80;
	blocks = (uint64_t*)calloc(padded_length/sizeof(uint64_t),sizeof(uint64_t));
	start = omp_get_wtime();
	#pragma omp parallel
	{
		#pragma omp sections private(my_blocks_read, i)
		{
			#pragma omp section
			{
				for(i = 0; i < message_length/8; i+=9)
				{
					my_blocks_read = fread((void*)(temp_read),sizeof(uint8_t),72,infile);
					if(i+9 >= message_length/8 && my_blocks_read < 72)
					{
						//printf("read: %d\n",my_blocks_read);
						temp_read[my_blocks_read] = 1;
						memset(temp_read+my_blocks_read+1,0,72-my_blocks_read-1);
					}
					#pragma omp critical(blocks_write)
					{
						memcpy(blocks+i,temp_read,72);
					}
					#pragma omp critical(blocks_count_update)
					blocks_read++;
					//printf("Message length = %d\nBytes to add = %d\nPadded length = %d\n",message_length,bytes_to_pad,padded_length);
				}
				#pragma omp critical(blocks_write)
				{
					/*for(i = message_length & 0xFFFFFFFFFFFFFFF8; i < message_length; i++)
					{
						padded_message[i] = blocks
					}*/
					//memcpy(padded_message+message_length-1,blocks+message_length/8-1,8);
					//padded_message[message_length] = 1;
					//padded_message[padded_length-1] |= 0x80;
					//memcpy(blocks+message_length/8,padded_message+message_length,bytes_to_pad);
					if(my_blocks_read == 72)
					blocks[message_length/8] = 1;
					//blocks[padded_length/8-1] |= 0x80;
					blocks[padded_length/8-1] = 0x8000000000000000;
				}
				#pragma omp critical(blocks_count_update)
				blocks_read += bytes_to_pad/8;
			}
			#pragma omp section
			{
				for(i = 0; i < padded_length/8; i+=9)
				{
					do
					{
						#pragma omp critical(blocks_count_update)
						my_blocks_read = blocks_read;
						//printf("Blocks read: %d\n",blocks_read);
					}
					while(my_blocks_read < 1);
					
					for(y = 0; y < 5; y++)
					{
						for(x = 0; x < 5; x++)
						{
							if(x+5*y < 576/64)
							{
								//#pragma omp critical(blocks_write)
								//{
									//printf("%.16lX\n",blocks[x+5*y+i]);
									S[x*STATE_ROWS+y] ^= blocks[x+5*y+i];
								//}
							}
						}
					}
					for(j = 0; j < 24; j++)
					{
						//θ step
						for(x = 0; x < STATE_ROWS; x++)
						{
							C[x] = S[x*STATE_ROWS] ^ S[x*STATE_ROWS+1] ^ S[x*STATE_ROWS+2] ^ S[x*STATE_ROWS+3] ^ S[x*STATE_ROWS+4];
						}
						for(x = 0; x < STATE_ROWS; x++)
						{
							D[x] = C[(x+4)%STATE_ROWS] ^ rot(C[(x+1)%STATE_ROWS],1);
						}
						for(x = 0; x < STATE_ROWS; x++)
						{
							for(y = 0; y < STATE_ROWS; y++)
							{
								S[x*STATE_ROWS+y] = S[x*STATE_ROWS+y] ^ D[x];
							}
						}

						//ρ and π steps
						for(x = 0; x < 5; x++)
						{
							for(y = 0; y < 5; y++)
							{
								B[y*STATE_ROWS+((2*x+3*y)%STATE_ROWS)] = rot(S[x*STATE_ROWS+y], rot_const[x][y]);
							}
						}

						//χ step
						//temp_time = omp_get_wtime();
						
						for(x = 0; x < 5; x++)
						{
							for(y = 0; y < 5; y++)
							{
								S[x*STATE_ROWS+y] = B[x*STATE_ROWS+y] ^ ((~B[((x+1)%STATE_ROWS)*STATE_ROWS+y]) & B[((x+2)%STATE_ROWS)*STATE_ROWS+y]);
							}
						}
						//profile_time += omp_get_wtime() - temp_time;

						//ι step
						S[0] = S[0] ^ RC[j];
					}
					#pragma omp critical(blocks_count_update)
					//#pragma omp atomic
					blocks_read--;
				}
			}
		}
	
	
	count = 0;
	
	/*for(y = 0; y < 5 && count < 8; y++)
	{
		for(x = 0; x < 5 && count < 8; x++)
		{
			if(x+5*y < 576/64)
			{
				sprintf(temp,"%.16lX",reverse(S[x*STATE_ROWS+y]));
				strcat(output,temp);
				count++;
			}
		}
	}*/
	#pragma omp for private(temp)
	for(x = 0; x < 8; x++)
	{
		sprintf(temp,"%.16lX",reverse(S[(5*x)%24]));
		memcpy(output+x*16,temp,16);
	}
	}
	//keccak_func(S);
	start = omp_get_wtime() - start;
	//printf("\n");
	/*for(i = 0; i < padded_length/8; i++)
	{
		printf("%.16lX\n",blocks[i]);
	}
	printf("\n");*/
	printf("%s %s\n",output,argv[1]);
	
	printf("Time: %lf\nProfiled part: %lf\n",start,profile_time);
	return 0;
}
